#include "../lib/basic.h"
#include <cstdio>
#include <cstring>
#include <cstdlib>
#include <iostream>
#include <string>
#include <vector>
#include <windows.h>
#include <cassert>
using namespace std;
using namespace Eigen;

const int thread_batch = 1000;
int seed;

struct kmeans_configuration
{
	int rf_size;
	int centroid_num;
	bool whitening;
	int iterations;
	int patch_n;
	int rand_seed;
	int thread_n;
	vector<string> files;
	string c_file, c_txt;
	bool pseudo_random() { return rand_seed == -1; }
};

bool parse_configuration(int argc, char *argv[], kmeans_configuration *config)
{
	config->rf_size = 10;
	config->centroid_num = 100;
	config->whitening = false;
	config->iterations = 50;
	config->thread_n = 16;
	config->rand_seed = -1;
	config->files.clear();
	config->c_file = "c.bin";
	config->c_txt = "";
	for (int i = 1; i < argc; ++i) {
		// common options
		int *p_val = NULL, lbound = 1;
		if (strcmp(argv[i], "-c") == 0)
			p_val = &config->centroid_num;
		if (strcmp(argv[i], "-k") == 0)
			p_val = &config->iterations;
		if (strcmp(argv[i], "-n") == 0)
			p_val = &config->patch_n;
		if (strcmp(argv[i], "-r") == 0)
			p_val = &config->rf_size;
		if (strcmp(argv[i], "-s") == 0) {
			p_val = &config->rand_seed;
			lbound = -1;
		}
		if (strcmp(argv[i], "-t") == 0)
			p_val = &config->thread_n;
		if (p_val != NULL) {
			if (i + 1 >= argc) {
				printf("error with parameters: no number after '%s'!\n", argv[i]);
				return false;
			}
			++i;
			*p_val = s_to_i(argv[i], -lbound - 1);
			if (*p_val < lbound) {
				printf("error with parameters: invalid number after '%s'!\n", argv[i]);
				return false;
			}
			continue;
		}
		// whitening option
		if (strcmp(argv[i], "-w") == 0) {
			config->whitening = true;
			continue;
		}
		// centroid file option
		if (strcmp(argv[i], "-cen") == 0) {
			if (i + 1 >= argc) {
				printf("error with parameters: no <c_file> after '%s'!\n", argv[i]);
				return false;
			}
			++i;
			config->c_file = string(argv[i]);
			continue;
		}
		// text file option
		if (strcmp(argv[i], "-txt") == 0) {
			if (i + 1 >= argc) {
				printf("error with parameters: no <c_txt> after '%s'!\n", argv[i]);
				return false;
			}
			++i;
			config->c_txt = string(argv[i]);
			continue;
		}
		// image file
		config->files.push_back(string(argv[i]));
	}
	return true;
}

void display_configuration(const kmeans_configuration &config)
{
	printf("Configuration:\n");
	printf("  common.centroid_num=%d\n", config.centroid_num);
	printf("  common.rf_size=%d\n", config.rf_size);
	printf("  common.whitening=%d\n", config.whitening);
	printf("  iterations=%d\n", config.iterations);
	printf("  patch_n=%d\n", config.patch_n);
	printf("  thread_n=%d\n", config.thread_n);
	printf("  rand_seed=%d\n", config.rand_seed);
	printf("  c_file='%s'\n", config.c_file.c_str());
	printf("  c_txt='%s'\n", config.c_txt.c_str());
	printf("  files.size()=%d\n", config.files.size());
	for (int i = 0; i < config.files.size(); ++i)
		printf("    files[%d]='%s'\n", i, config.files[i].c_str());
	printf("\n");
}

struct whitening_data
{
	bool enabled;
	float *M;
	float **P;
};

whitening_data run_whitening(MatrixR &patch)
{
	int patch_n = patch.rows();
	int patch_size = patch.cols();

	printf("  - computing covariance...\n");
	RowVectorR mean = patch.colwise().mean();
	MatrixR centered = patch.rowwise() - mean;
	MatrixR cov = centered.adjoint() * centered / (patch_n - 1);

	printf("  - computing eigenvectors...\n");
	EigenSolver<MatrixR> rw_solver(cov);
	MatrixC D0 = rw_solver.eigenvalues().asDiagonal();
	MatrixC V0 = rw_solver.eigenvectors();

	MatrixR D(patch_size, patch_size);
	MatrixR V(patch_size, patch_size);
	for (int i = 0; i < patch_size; ++i)
		for (int j = 0; j < patch_size; ++j) {
			if (i == j) {
				real_num d = D0(i, j).real();
				D(i, j) = sqrt(1./ (d + 0.1));
			} else {
				D(i, j) = 0.;
			}
			V(i, j) = V0(i, j).real();
		}
		
	printf("  - computing P...\n");
	MatrixR P = V * D * V.transpose();

	printf("  - modifying patches...\n");
	patch = (patch.rowwise() - mean) * P;

	printf("  - saving data...\n");
	whitening_data data;
	data.enabled = true;
	data.M = new float[patch_size];
	for (int j = 0; j < patch_size; ++j)
		data.M[j] = (float)mean(j);
	data.P = new float*[patch_size];
	for (int i = 0; i < patch_size; ++i) {
		data.P[i] = new float[patch_size];
		for (int j = 0; j < patch_size; ++j)
			data.P[i][j] = (float)P(i, j);
	}
	return data;
}

struct multiply_data
{
	bool flag;
	MatrixR a, b, c;
	//int x, y, z;
};

DWORD WINAPI multiply_task(LPVOID lpParam) // c = a * b^T
{
	multiply_data *data = (multiply_data*)lpParam;
	//*
	if (data->flag) {
		//data->z = data->x * data->y;
		data->c = data->a * data->b.transpose();
		//printf("(%d,%d) * (%d,%d) -> (%d,%d)\n",data->a.rows(),data->a.cols(), data->b.rows(),data->b.cols(), data->c.rows(),data->c.cols());
	}//*/
	return TRUE;
}

void kmeans_iteration(const MatrixR &patch, MatrixR &centroid, int thread_n)
{
	const int max_thread_n = 32;
	static multiply_data data[max_thread_n];
	static HANDLE threads[max_thread_n];
	thread_n = min(thread_n, max_thread_n);

	int centroid_num = (int)centroid.rows();
	int patch_size = (int)centroid.cols();
	int patch_n = (int)patch.rows();
	VectorR c2 = centroid.cwiseProduct(centroid).rowwise().sum() / 2.;
	MatrixR sum(centroid_num, patch_size);
	int *cnt = new int[centroid_num];
	sum.setZero();
	memset(cnt, 0, centroid_num * sizeof(int));
	for (int cur = 0; cur < patch_n; cur += thread_batch * thread_n) {
		for (int i = 0; i < thread_n; ++i) {
			int st = cur + i * thread_batch;
			int en = min(st + thread_batch, patch_n);
			if (st < en) {
				data[i].flag = true;
				data[i].a = centroid;
				data[i].b = patch.block(st, 0, en-st, patch_size);
			} else
				data[i].flag = false;
			DWORD t_id;
			threads[i] = CreateThread(NULL, 0, (LPTHREAD_START_ROUTINE)multiply_task, (LPVOID)&data[i], 0, &t_id);
			if (threads[i] == NULL) {
				printf(" - unable to create thread!!! error=%d\n", GetLastError());
				exit(1);
			}
		}
		WaitForMultipleObjects(thread_n, threads, TRUE, INFINITE);
		for (int i = 0; i < thread_n; ++i)
			CloseHandle(threads[i]);
		for (int i = 0; i < thread_n; ++i) {
			if (!data[i].flag)
				continue;
			data[i].c = data[i].c.colwise() - c2;
			for (int j = 0; j < data[i].b.rows(); ++j) {
				int ind;
				data[i].c.col(j).maxCoeff(&ind);
				sum.row(ind) += data[i].b.row(j);
				++cnt[ind];
			}
		}
	}
	for (int i = 0; i < centroid_num; ++i)
		centroid.row(i) = sum.row(i) / (real_num)max(cnt[i], 1);
	delete cnt;
}

void save_centroids(MatrixR centroid, whitening_data w, int rf_size, const char *c_file, const char *c_txt)
{
	int centroid_num = (int)centroid.rows();
	int patch_size = (int)centroid.cols();
	FILE *f = fopen(c_file, "wb");
	if (f == NULL) {
		printf("cannot save centroids to %s!\n", c_file);
		exit(1);
	}
	fwrite(&rf_size, 4, 1, f);
	fwrite(&centroid_num, 4, 1, f);
	fwrite(&patch_size, 4, 1, f);
	float *buf = new float[patch_size];
	for (int i = 0; i < centroid_num; ++i) {
		for (int j = 0; j < patch_size; ++j)
			buf[j] = centroid(i, j);
		fwrite(buf, sizeof(float) * patch_size, 1, f);
	}
	delete buf;
	int wh = w.enabled ? 1 : 0;
	fwrite(&wh, 4, 1, f);
	if (wh == 1) {
		fwrite(w.M, sizeof(float) * patch_size, 1, f);
		for (int i = 0 ; i < patch_size; ++i)
			fwrite(w.P[i], sizeof(float) * patch_size, 1, f);
	}
	fclose(f);

	if (strcmp(c_txt, "") != 0) {
		f = fopen(c_txt, "w");
		if (f == NULL) {
			printf("cannot save centroids to text file %s!\n", c_txt);
			exit(1);
		}
		fprintf(f, "%d %d\n", centroid_num, patch_size);
		for (int i = 0; i < centroid_num; ++i) {
			for (int j = 0; j < patch_size; ++j)
				fprintf(f, "%.4lf ", (double)centroid(i, j));
			fprintf(f, "\n");
		}
		fclose(f);
	}
}

int main(int argc, char *argv[])
{
	_timer tot_time;

	kmeans_configuration config;
	if (!parse_configuration(argc, argv, &config))
		return 1;
	display_configuration(config);
	if (!config.pseudo_random())
		srand(config.rand_seed);

	image_data_set img;
	for (int i = 0; i < (int)config.files.size(); ++i) {
		printf("loading image %d/%d (%s) ...\n", i + 1, config.files.size(), config.files[i].c_str());
		int cnt = img.load_bin(config.files[i].c_str());
		printf("loaded %d images\n", cnt);
	}
	printf("totally loaded %d images.\n", img.size());
	if (img.size() == 0) {
		printf("no image!\n");
		return 1;
	}

	int channels = img[0].channels;
	for (int i = 0; i < (int)img.size(); ++i)
		if (img[i].channels != channels) {
			printf("wrong number of image channels (c[1]=%d c[%d]=%d)!\n", channels, i+1, img[i].channels);
			return 1;
		}
	int rf_size = config.rf_size;
	int patch_size = rf_size * rf_size * channels;
	int centroid_num = config.centroid_num;
	int thread_n = config.thread_n;
	int patch_n = config.patch_n;
	
	printf("sample patches...\n");
	MatrixR patch(config.patch_n, patch_size);
	seed = 123456;
	for (int cur = 0; cur < config.patch_n; ++cur) {
		int i, x, y;
		if (config.pseudo_random()) {
			i = cur % img.size();
			x = seed % (img[i].height - config.rf_size + 1);
			seed = (seed * 17 + 11) % 100007;
			y = seed % (img[i].width - config.rf_size + 1);
			seed = (seed * 17 + 11) % 100007;
		} else {
			i = random_int() % img.size();
			x = random_int() % (img[i].height - config.rf_size + 1);
			y = random_int() % (img[i].width - config.rf_size + 1);
		}
		real_num *p = new real_num[patch_size];
		extract_patch(img[i], x, y, rf_size, p);
		normalize_patch(p, patch_size);
		for (int j = 0; j < patch_size; ++j)
			patch(cur, j) = p[j];
		delete p;
	}

	whitening_data w_data;
	if (config.whitening) {
		printf("whitening...\n");
		w_data = run_whitening(patch);
	} else
		w_data.enabled = false;

	printf("computing centroids...\n");
	MatrixR centroid(centroid_num, patch_size);
	seed = 123456;
	for (int i = 0; i < centroid_num; ++i) {
		int k;
		if (config.pseudo_random()) {
			k = seed % patch_n;
			seed = (seed * 19 + 11) % 100007;
		} else {
			k = random_int() % patch_n;
		}
		//printf("i=%d k=%d\n", i, k);
		assert(i >= 0 && i < (int)centroid.rows());
		assert(k >= 0 && k < (int)patch.rows());
		centroid.row(i) = patch.row(k);
	}
	_timer timer(config.iterations);
	for (int itr = 0; itr < config.iterations; ++itr) {
		printf("kmeans iteration %d/%d [reamin %s]...\n", itr, config.iterations, timer.predict(itr));
		kmeans_iteration(patch, centroid, config.thread_n);
	}

	printf("saving centroids...\n");
	save_centroids(centroid, w_data, rf_size, config.c_file.c_str(), config.c_txt.c_str());
	
	//freopen("t.txt", "w", stdout);

	printf("done. [time=%s]\n", tot_time.count());
	freopen("__kmeans_time.txt", "w", stdout);
	printf("done. [time=%s]\n", tot_time.count());
	
	return 0;
}
